package cn.xzc.job.cocurrent;

import cn.xzc.job.ErrorHandler;
import cn.xzc.job.JobScheduler;
import cn.xzc.job.Trigger;
import cn.xzc.job.pojo.ConcurrentReferenceHashMap;
import cn.xzc.job.util.JobUtils;

import java.time.Clock;
import java.util.Date;
import java.util.Map;
import java.util.concurrent.*;

@SuppressWarnings("serial")
public class ThreadPoolJobScheduler extends ExecutorConfigurationSupport
        implements AsyncListenableJobExecutor, SchedulingJobExecutor, JobScheduler {
    // Underlying ScheduledFutureTask to user-level ListenableFuture handle, if any
    private final Map<Object, ListenableFuture<?>> listenableFutureMap =
            new ConcurrentReferenceHashMap<>(16, ConcurrentReferenceHashMap.ReferenceType.WEAK);
    private volatile int poolSize = 1;
    private volatile boolean removeOnCancelPolicy;
    private volatile boolean continueExistingPeriodicTasksAfterShutdownPolicy;
    private volatile boolean executeExistingDelayedTasksAfterShutdownPolicy = true;
    
    private volatile ErrorHandler errorHandler;
    private Clock clock = Clock.systemDefaultZone();
    
    private ScheduledExecutorService scheduledExecutor;

    /**
     * Set whether to continue existing periodic tasks even when this executor has been shutdown.
     * <p>Default is {@code false}. If set to {@code true}, the target executor will be
     * switched into continuing periodic tasks (if possible).
     * <p><b>This setting can be modified at runtime, for example through JMX.</b>
     *
     * @see ScheduledThreadPoolExecutor#setContinueExistingPeriodicTasksAfterShutdownPolicy
     */
    public void setContinueExistingPeriodicTasksAfterShutdownPolicy(boolean flag) {
        if (this.scheduledExecutor instanceof ScheduledThreadPoolExecutor) {
            ((ScheduledThreadPoolExecutor) this.scheduledExecutor).setContinueExistingPeriodicTasksAfterShutdownPolicy(flag);
        }
        this.continueExistingPeriodicTasksAfterShutdownPolicy = flag;
    }

    /**
     * Set whether to execute existing delayed tasks even when this executor has been shutdown.
     * <p>Default is {@code true}. If set to {@code false}, the target executor will be
     * switched into dropping remaining tasks (if possible).
     * <p><b>This setting can be modified at runtime, for example through JMX.</b>
     *
     * @see ScheduledThreadPoolExecutor#setExecuteExistingDelayedTasksAfterShutdownPolicy
     */
    public void setExecuteExistingDelayedTasksAfterShutdownPolicy(boolean flag) {
        if (this.scheduledExecutor instanceof ScheduledThreadPoolExecutor) {
            ((ScheduledThreadPoolExecutor) this.scheduledExecutor).setExecuteExistingDelayedTasksAfterShutdownPolicy(flag);
        }
        this.executeExistingDelayedTasksAfterShutdownPolicy = flag;
    }

    /**
     * Set a custom {@link ErrorHandler} strategy.
     */
    public void setErrorHandler(ErrorHandler errorHandler) {
        this.errorHandler = errorHandler;
    }

    @Override
    public Clock getClock() {
        return this.clock;
    }

    /**
     * Set the clock to use for scheduling purposes.
     * <p>The default clock is the system clock for the default time zone.
     *
     * @see Clock#systemDefaultZone()
     */
    public void setClock(Clock clock) {
        this.clock = clock;
    }

    @Override
    protected ExecutorService initializeExecutor(
            ThreadFactory threadFactory, RejectedExecutionHandler rejectedExecutionHandler) {

        this.scheduledExecutor = createExecutor(this.poolSize, threadFactory, rejectedExecutionHandler);

        if (this.scheduledExecutor instanceof ScheduledThreadPoolExecutor) {
            ScheduledThreadPoolExecutor scheduledPoolExecutor = (ScheduledThreadPoolExecutor) this.scheduledExecutor;
            if (this.removeOnCancelPolicy) {
                scheduledPoolExecutor.setRemoveOnCancelPolicy(true);
            }
            if (this.continueExistingPeriodicTasksAfterShutdownPolicy) {
                scheduledPoolExecutor.setContinueExistingPeriodicTasksAfterShutdownPolicy(true);
            }
            if (!this.executeExistingDelayedTasksAfterShutdownPolicy) {
                scheduledPoolExecutor.setExecuteExistingDelayedTasksAfterShutdownPolicy(false);
            }
        }

        return this.scheduledExecutor;
    }

    /**
     * Create a new {@link ScheduledExecutorService} instance.
     * <p>The default implementation creates a {@link ScheduledThreadPoolExecutor}.
     * Can be overridden in subclasses to provide custom {@link ScheduledExecutorService} instances.
     *
     * @param poolSize                 the specified pool size
     * @param threadFactory            the ThreadFactory to use
     * @param rejectedExecutionHandler the RejectedExecutionHandler to use
     * @return a new ScheduledExecutorService instance
     * @see #init()
     * @see java.util.concurrent.ScheduledThreadPoolExecutor
     */
    protected ScheduledExecutorService createExecutor(
            int poolSize, ThreadFactory threadFactory, RejectedExecutionHandler rejectedExecutionHandler) {

        return new ScheduledThreadPoolExecutor(poolSize, threadFactory, rejectedExecutionHandler);
    }

    /**
     * Return the underlying ScheduledExecutorService for native access.
     *
     * @return the underlying ScheduledExecutorService (never {@code null})
     * @throws IllegalStateException if the ThreadPoolJobScheduler hasn't been initialized yet
     */
    public ScheduledExecutorService getScheduledExecutor() throws IllegalStateException {
        return this.scheduledExecutor;
    }

    /**
     * Return the underlying ScheduledThreadPoolExecutor, if available.
     *
     * @return the underlying ScheduledExecutorService (never {@code null})
     * @throws IllegalStateException if the ThreadPoolJobScheduler hasn't been initialized yet
     *                               or if the underlying ScheduledExecutorService isn't a ScheduledThreadPoolExecutor
     * @see #getScheduledExecutor()
     */
    public ScheduledThreadPoolExecutor getScheduledThreadPoolExecutor() throws IllegalStateException {
        return (ScheduledThreadPoolExecutor) this.scheduledExecutor;
    }

    /**
     * Return the current pool size.
     * <p>Requires an underlying {@link ScheduledThreadPoolExecutor}.
     *
     * @see #getScheduledThreadPoolExecutor()
     * @see java.util.concurrent.ScheduledThreadPoolExecutor#getPoolSize()
     */
    public int getPoolSize() {
        if (this.scheduledExecutor == null) {
            // Not initialized yet: assume initial pool size.
            return this.poolSize;
        }
        return getScheduledThreadPoolExecutor().getPoolSize();
    }

    /**
     * Set the ScheduledExecutorService's pool size.
     * Default is 1.
     * <p><b>This setting can be modified at runtime, for example through JMX.</b>
     */
    public void setPoolSize(int poolSize) {
        if (this.scheduledExecutor instanceof ScheduledThreadPoolExecutor) {
            ((ScheduledThreadPoolExecutor) this.scheduledExecutor).setCorePoolSize(poolSize);
        }
        this.poolSize = poolSize;
    }

    /**
     * Return the number of currently active threads.
     * <p>Requires an underlying {@link ScheduledThreadPoolExecutor}.
     *
     * @see #getScheduledThreadPoolExecutor()
     * @see java.util.concurrent.ScheduledThreadPoolExecutor#getActiveCount()
     */
    public int getActiveCount() {
        if (this.scheduledExecutor == null) {
            // Not initialized yet: assume no active threads.
            return 0;
        }
        return getScheduledThreadPoolExecutor().getActiveCount();
    }


    /**
     * Set the remove-on-cancel mode on {@link ScheduledThreadPoolExecutor}.
     * <p>Default is {@code false}. If set to {@code true}, the target executor will be
     * switched into remove-on-cancel mode (if possible).
     * <p><b>This setting can be modified at runtime, for example through JMX.</b>
     *
     * @see ScheduledThreadPoolExecutor#setRemoveOnCancelPolicy
     */
    public void setRemoveOnCancelPolicy(boolean flag) {
        if (this.scheduledExecutor instanceof ScheduledThreadPoolExecutor) {
            ((ScheduledThreadPoolExecutor) this.scheduledExecutor).setRemoveOnCancelPolicy(flag);
        }
        this.removeOnCancelPolicy = flag;
    }


    // SchedulingTaskExecutor implementation

    @Override
    public void execute(Runnable task) {
        Executor executor = getScheduledExecutor();
        try {
            executor.execute(errorHandlingTask(task, false));
        } catch (RejectedExecutionException ex) {
            throw new RejectedExecutionException("Executor [" + executor + "] did not accept task: " + task, ex);
        }
    }

    @Override
    public Future<?> submit(Runnable task) {
        ExecutorService executor = getScheduledExecutor();
        try {
            return executor.submit(errorHandlingTask(task, false));
        } catch (RejectedExecutionException ex) {
            throw new RejectedExecutionException("Executor [" + executor + "] did not accept task: " + task, ex);
        }
    }

    @Override
    public <T> Future<T> submit(Callable<T> task) {
        ExecutorService executor = getScheduledExecutor();
        try {
            Callable<T> taskToUse = task;
            ErrorHandler errorHandler = this.errorHandler;
            if (errorHandler != null) {
                taskToUse = new DelegatingErrorHandlingCallable<>(task, errorHandler);
            }
            return executor.submit(taskToUse);
        } catch (RejectedExecutionException ex) {
            throw new RejectedExecutionException("Executor [" + executor + "] did not accept task: " + task, ex);
        }
    }

    @Override
    public ListenableFuture<?> submitListenable(Runnable task) {
        ExecutorService executor = getScheduledExecutor();
        try {
            ListenableFutureTask<Object> listenableFuture = new ListenableFutureTask<>(task, null);
            executeAndTrack(executor, listenableFuture);
            return listenableFuture;
        } catch (RejectedExecutionException ex) {
            throw new RejectedExecutionException("Executor [" + executor + "] did not accept task: " + task, ex);
        }
    }

    @Override
    public <T> ListenableFuture<T> submitListenable(Callable<T> task) {
        ExecutorService executor = getScheduledExecutor();
        try {
            ListenableFutureTask<T> listenableFuture = new ListenableFutureTask<>(task);
            executeAndTrack(executor, listenableFuture);
            return listenableFuture;
        } catch (RejectedExecutionException ex) {
            throw new RejectedExecutionException("Executor [" + executor + "] did not accept task: " + task, ex);
        }
    }

    private void executeAndTrack(ExecutorService executor, ListenableFutureTask<?> listenableFuture) {
        Future<?> scheduledFuture = executor.submit(errorHandlingTask(listenableFuture, false));
        this.listenableFutureMap.put(scheduledFuture, listenableFuture);
        listenableFuture.addCallback(result -> this.listenableFutureMap.remove(scheduledFuture),
                ex -> this.listenableFutureMap.remove(scheduledFuture));
    }

    /**
     * 强制关闭runnable
     *
     * @param task the task to cancel (typically a {@link RunnableFuture}) 任务
     */
    @Override
    protected void cancelRemainingTask(Runnable task) {
        super.cancelRemainingTask(task);
        // Cancel associated user-level ListenableFuture handle as well
        ListenableFuture<?> listenableFuture = this.listenableFutureMap.get(task);
        if (listenableFuture != null) {
            listenableFuture.cancel(true);
        }
    }


    // TaskScheduler implementation

    @Override
    
    public ScheduledFuture<?> schedule(Runnable task, Trigger trigger) {
        ScheduledExecutorService executor = getScheduledExecutor();
        try {
            ErrorHandler errorHandler = this.errorHandler;
            if (errorHandler == null) {
                //需要重复执行的任务的默认异常处理策略
                errorHandler = JobUtils.getDefaultErrorHandler(true);
            }
            return new ReschedulingRunnable(task, trigger, this.clock, executor, errorHandler).schedule();
        } catch (RejectedExecutionException ex) {
            throw new RejectedExecutionException("Executor [" + executor + "] did not accept task: " + task, ex);
        }
    }

    @Override
    public ScheduledFuture<?> schedule(Runnable task, Date startTime) {
        ScheduledExecutorService executor = getScheduledExecutor();
        long initialDelay = startTime.getTime() - this.clock.millis();
        try {
            return executor.schedule(errorHandlingTask(task, false), initialDelay, TimeUnit.MILLISECONDS);
        } catch (RejectedExecutionException ex) {
            throw new RejectedExecutionException("Executor [" + executor + "] did not accept task: " + task, ex);
        }
    }

    @Override
    public ScheduledFuture<?> scheduleAtFixedRate(Runnable task, Date startTime, long period) {
        ScheduledExecutorService executor = getScheduledExecutor();
        long initialDelay = startTime.getTime() - this.clock.millis();
        try {
            return executor.scheduleAtFixedRate(errorHandlingTask(task, true), initialDelay, period, TimeUnit.MILLISECONDS);
        } catch (RejectedExecutionException ex) {
            throw new RejectedExecutionException("Executor [" + executor + "] did not accept task: " + task, ex);
        }
    }

    @Override
    public ScheduledFuture<?> scheduleAtFixedRate(Runnable task, long period) {
        ScheduledExecutorService executor = getScheduledExecutor();
        try {
            return executor.scheduleAtFixedRate(errorHandlingTask(task, true), 0, period, TimeUnit.MILLISECONDS);
        } catch (RejectedExecutionException ex) {
            throw new RejectedExecutionException("Executor [" + executor + "] did not accept task: " + task, ex);
        }
    }

    @Override
    public ScheduledFuture<?> scheduleWithFixedDelay(Runnable task, Date startTime, long delay) {
        ScheduledExecutorService executor = getScheduledExecutor();
        long initialDelay = startTime.getTime() - this.clock.millis();
        try {
            return executor.scheduleWithFixedDelay(errorHandlingTask(task, true), initialDelay, delay, TimeUnit.MILLISECONDS);
        } catch (RejectedExecutionException ex) {
            throw new RejectedExecutionException("Executor [" + executor + "] did not accept task: " + task, ex);
        }
    }

    @Override
    public ScheduledFuture<?> scheduleWithFixedDelay(Runnable task, long delay) {
        ScheduledExecutorService executor = getScheduledExecutor();
        try {
            return executor.scheduleWithFixedDelay(errorHandlingTask(task, true), 0, delay, TimeUnit.MILLISECONDS);
        } catch (RejectedExecutionException ex) {
            throw new RejectedExecutionException("Executor [" + executor + "] did not accept task: " + task, ex);
        }
    }


    private Runnable errorHandlingTask(Runnable task, boolean isRepeatingTask) {
        return JobUtils.decorateTaskWithErrorHandler(task, this.errorHandler, isRepeatingTask);
    }


    private static class DelegatingErrorHandlingCallable<V> implements Callable<V> {

        private final Callable<V> delegate;

        private final ErrorHandler errorHandler;

        public DelegatingErrorHandlingCallable(Callable<V> delegate, ErrorHandler errorHandler) {
            this.delegate = delegate;
            this.errorHandler = errorHandler;
        }

        @Override
        
        public V call() throws Exception {
            try {
                return this.delegate.call();
            } catch (Throwable ex) {
                this.errorHandler.handleError(ex);
                return null;
            }
        }
    }

}
